/*
 * Optimizer.cpp
 *
 *      Author: Jerome RD Soine and Christoph A Brand
 *      Institution: Schwarz goup, Institute for Theoretical Physics, Heidelberg University, Germany
 */

#include "../../include/core/Optimizer.h"
#include "../../include/core/utility/VectorOperations.h"
#include <cmath>
#include <unistd.h>
#include <iostream>

Optimizer::Optimizer(CellModel * CELL_MODEL, SubstrateModel * SUBSTRATE_MODEL, Comparator * COMPARATOR, string SOLUTION_FOLDER, unsigned int N_ITERATIONS, unsigned int N_THREADS, double TOLERANCE, double GRADIENT_SPACING_BASE){
	cell_model = CELL_MODEL;
	substrate_model = SUBSTRATE_MODEL;
	comparator = COMPARATOR;
	n_iterations = N_ITERATIONS;
	n_threads = N_THREADS;
	tolerance = TOLERANCE;
	gradient_spacing_base = GRADIENT_SPACING_BASE;
	WorkerManager::initialize(cell_model, substrate_model, comparator, N_THREADS);
	worker_manager = WorkerManager::getInstance();
	solution_folder = SOLUTION_FOLDER;
	solution = CELL_MODEL->get_parameters();

	int deg_of_freedom = cell_model->n_parameters();
	set_gradient_spacing_scaling(vector<double>(deg_of_freedom,1));
}

vector<double> Optimizer::gradient(vector<double> parameters){

    vector <double> df (parameters.size(),0);
    double f;
    vector<vector<double> > dx_list (parameters.size());
    vector<bool> joined_list (parameters.size(),false);

    vector<Worker*> workers;
    for(unsigned int i=0;i<n_threads;i++)
    	workers.push_back(worker_manager->getWorker(parameters,-1));

    // Calculate function value and seed configuration
    workers[0]->init(parameters,-1);
    workers[0]->work();
    f = workers[0]->get_fitness();
    vector<double> seed = workers[0]->get_final_configuration();

    for(unsigned int i=0;i<parameters.size();i++)
    {
    	cout << "Parameter " << i << " has value " << parameters[i] << " and will be displaced by " << gradient_spacing_base <<" * " << gradient_spacing_scaling[i] << " = " << gradient_spacing_base * gradient_spacing_scaling[i]<< endl;
    	dx_list[i] = parameters;
    	dx_list[i][i] += gradient_spacing_base * gradient_spacing_scaling[i];
    }

    unsigned int counter = 0;

    vector<double> thread_to_component (n_threads,0);

    if(parameters.size() <= n_threads){
    	for(unsigned int i=0; i<parameters.size(); i++){
    		workers[i]->init(dx_list[i],seed,-1);
    		workers[i]->start();
    	}
    	for(unsigned int i=0; i<parameters.size(); i++){
    		workers[i]->join();
    	}
    	for(unsigned int i=0; i<parameters.size(); i++){
    		df[i] = workers[i]->get_fitness();
    	}
    } else {
    	for(unsigned int i=0; i<n_threads; i++){
    		thread_to_component[i] = counter;
    		workers[i]->init(dx_list[counter],seed,-1);
    		workers[i]->start();
    		counter++;
    	}
    	while(true){
    		for(unsigned int i=0; i<n_threads; i++){
    				bool joined = workers[i]->try_join_for(10l);
    				if(joined){
    					df[thread_to_component[i]] = workers[i]->get_fitness();
    					joined_list[thread_to_component[i]] = true;
    					if(counter < parameters.size()){
    						workers[i]->init(dx_list[counter],seed,-1);
    						thread_to_component[i]=counter;
    						workers[i]->start();
    						counter++;
    					} else {

    					}
    				}
    			}
    		bool all_joined = true;
    		for(unsigned int j=0;j<parameters.size();j++)
    		    if(!joined_list[j])
    		    	all_joined = false;
    		if(all_joined)
    			break;
    	}
    }

    for(unsigned int i=0;i<parameters.size();i++)
    {
        df[i]=(df[i]-f)/ (gradient_spacing_base * gradient_spacing_scaling[i]);
    }

	for(unsigned int i=0;i<workers.size();i++)
		workers[i]->unlock();

	cout << "gradient = ";
			for(unsigned int i=0;i<df.size();i++)
				cout << df[i] << "   ";
	cout << "\n";

    return df;
}

vector<double> Optimizer::bracket_line_minimum(vector<double> start_pos, vector<double> direction, vector<double>& seed){
	cout << "Starting bracketing...\n";
	double eA, eC;

	const double gold = 1.618034;
	double initial_displacement = gradient_spacing_base*10;

	Worker * preWorker = worker_manager->getWorker(start_pos,-1);
	preWorker->start();

	vector<double> new_pos = start_pos;
	VectorOperations::add_to_std_vector(initial_displacement, direction, new_pos);
	Worker * preWorker2 = worker_manager->getWorker(new_pos,-1);
	preWorker2->start();

	preWorker->join();
	preWorker2->join();

	eA = preWorker->get_fitness();
	eC = preWorker2->get_fitness();

	int i = 0;
	if (eC >= eA) {
		while (eC > eA && i < 10) {
			initial_displacement *= 0.1;
			new_pos = start_pos;
			VectorOperations::add_to_std_vector(initial_displacement, direction, new_pos);
			preWorker2->unlock();
			preWorker2->init(new_pos, preWorker->get_final_configuration(),-1);
			preWorker2->work();
			eC = preWorker2->get_fitness();
			i++;
		}
	}
	if (eC >= eA) {
		preWorker->unlock();
		preWorker2->unlock();
		throw -1;
	}

	preWorker->unlock();
	preWorker2->unlock();

	vector<Worker*> worker;
	vector<double> pos_list;
	vector<double> l2_values_list;

	for(unsigned int i=0; i<n_threads; i++){
		double d = pow(gold,(double)(pos_list.size()+1)) * initial_displacement;
		pos_list.push_back(d);
		new_pos = start_pos;
		VectorOperations::add_to_std_vector(d, direction, new_pos);
		worker.push_back(worker_manager->getWorker(new_pos, preWorker2->get_final_configuration(),pos_list.size()-1));
		worker[i]->start();
		l2_values_list.push_back(-1.0);
	}

	double minimum_bracketed = false;

	while(!minimum_bracketed){
		for(unsigned int i=0; i<n_threads; i++){
			if(minimum_bracketed) break;
			bool joined = worker[i]->try_join_for(10l);
			if(joined){
				int id = worker[i]->get_id();
				double l2 = worker[i]->get_fitness();
				l2_values_list[id] = l2;
				seed = worker[i]->get_final_configuration();
//				worker[i]->unlock();
				if(l2_values_list[id]>=eC){
					minimum_bracketed = true;
				}
				if(!minimum_bracketed){
					for(int i=0; i<id; i++){
						if(l2_values_list[i] > 0 && l2_values_list[i] <= l2){
							minimum_bracketed = true;
							break;
						}
					}
				}
				if(!minimum_bracketed){
					//Check if smaller than any later values
					for(unsigned int i=id+1; i<l2_values_list.size(); i++){
						if(l2_values_list[i] > 0 && l2_values_list[i] >= l2){
							minimum_bracketed = true;
							break;
						}
					}
				}

				if(!minimum_bracketed){
					//Now we are sure we have to continue the search and want to create a new worker
					double d = pow(gold,(double)(pos_list.size()+1)) * initial_displacement;
					pos_list.push_back(d);
					new_pos = start_pos;
					VectorOperations::add_to_std_vector(d, direction, new_pos);
					worker[i]->init(new_pos, seed,pos_list.size()-1);
					l2_values_list.push_back(-1.0);
					worker[i]->start();
				}

			}
		}
	}

	for(unsigned int i=0; i<n_threads; i++){
		{
			worker[i]->join();
			int id = worker[i]->get_id();
			double l2 = worker[i]->get_fitness();
			l2_values_list[id] = l2;
		}
	}


	vector<double> all_l2s(l2_values_list.size()+2,-1);
	vector<double> all_pos(pos_list.size()+2,0);
	all_l2s[0] = eA;
	all_l2s[1] = eC;
	all_pos[0] = 0;
	all_pos[1] = initial_displacement;
	for(unsigned int i = 0; i<l2_values_list.size(); i++){
		all_l2s[i+2] = l2_values_list[i];
		all_pos[i+2] = pos_list[i];
	}

	int before_min_id = 0, min_id = 1, after_min_id=-1;
	for(unsigned int i=2; i<all_l2s.size();i++){
		double cl2 = all_l2s[i];
		if(cl2 >= 0){
			if(cl2 >= all_l2s[min_id]){
				after_min_id = i;
				break;
			} else {
				before_min_id = min_id;
				min_id = i;
			}
		}
	}

	if(after_min_id==-1){
		cout << "COULD NOT BRACKET MINIMUM\n";
		for(unsigned int i=0; i<l2_values_list.size();i++){
			cout << l2_values_list[i] << " - ";
		}
		cout << "\n";
		exit(0);
	}

	vector<double> ret;
	ret.push_back(all_pos[before_min_id]);
	ret.push_back(all_pos[min_id]);
	ret.push_back(all_pos[after_min_id]);
	ret.push_back(all_l2s[before_min_id]);
	ret.push_back(all_l2s[min_id]);
	ret.push_back(all_l2s[after_min_id]);

	if(min_id == 1) seed = preWorker2->get_final_configuration();
	else seed = worker[0]->get_final_configuration();

	for(unsigned int i=0; i<worker.size(); i++){
		worker[i]->unlock();
	}
	preWorker->unlock();
	preWorker2->unlock();


	return ret;
}

vector<double> Optimizer::line_minimization(vector<double> start_pos, vector<double> direction){
	cout << "Starting minimum search with tolerance " << tolerance <<"\n";
	bool no_direction_specified = true;
	for(unsigned int i=0; i<direction.size(); i++){
		if(abs(direction[i]) > 0 || std::isnan(direction[i])){
			no_direction_specified = false;
			break;
		}
	}
	if(no_direction_specified){
		cout << "No search direction specified (gradient is zero). Returning start_pos." << endl;
		return start_pos;
	}

	//Initialize
	const vector<double> orig_start_pos = start_pos;
	vector<double> ret;

	double before_min_displacement, min_displacement, after_min_displacement;
	VectorOperations::normalize_vector(direction);
	VectorOperations::change_sign_of_std_vector(direction);
	vector<double> seed;

	//Bracket minimum and make associated variables accessable
	vector<double> bracket;
	try{
		bracket = bracket_line_minimum(start_pos, direction, seed);
	}catch(int errorcode){
		//Could not bracket minimum - L2 value increases along given line.
		cout << "\n\n FATAL ERROR: COULD NOT BRACKET MINIMUM. RETURNING CURRENT POSITION AS MINIMUM AND CONTINUING.\n\n";
		return orig_start_pos;
	}
	cout << "Done bracketing minimum." << endl;

	VectorOperations::cout_std_vector(bracket);

	before_min_displacement = bracket[0];
	min_displacement = bracket[1];
	after_min_displacement = bracket[2];
	double value_at_before_min_displacement = bracket[3];
	double value_at_min_displacement = bracket[4];
	double value_at_after_min_displacement = bracket[5];
	ret = orig_start_pos;
	VectorOperations::add_to_std_vector(min_displacement, direction, ret);

	vector<Worker*> workers;
	for(unsigned int i=0;i<n_threads;i++)
		workers.push_back(worker_manager->getWorker(orig_start_pos,seed,i));

	while(abs(after_min_displacement - before_min_displacement) > tolerance){

		//Current position is always where the first of the three minimum bracketing points is.
		vector<double> current_pos = orig_start_pos;
		VectorOperations::add_to_std_vector(before_min_displacement, direction, current_pos);

		//Initialize positions, where the N_THREADS new values are to be calculated with equal spacing along the given interval. Integrate the three bracketing points.
		double dist = after_min_displacement - before_min_displacement;
		double step_size = dist/(n_threads+1);
		vector<double> values(n_threads+3,0);
		values[0] = value_at_before_min_displacement;
		values[values.size()-1] = value_at_after_min_displacement;
		unsigned int pos_min_displacement = (int) ceil((min_displacement-before_min_displacement)/step_size);
		values[pos_min_displacement] = value_at_min_displacement;

		//Calculate distances between points to which values either exist or are to be calculated.
   		vector<double> step_sizes(n_threads+2,0);
   		for(unsigned int i=0; i<values.size()-1; i++){
   			if(i<pos_min_displacement-1) step_sizes[i] = step_size;
   			else if(i+1==pos_min_displacement) step_sizes[i] = (min_displacement - before_min_displacement) - i * step_size;
   			else if(i==pos_min_displacement) step_sizes[i] = i*step_size - (min_displacement - before_min_displacement);
   			else step_sizes[i] = step_size;
   		}

   		//Save absolute positions in parameter space
   		vector<double> temp_pos=current_pos;
   		vector< vector<double> > positions;
   		positions.push_back(current_pos);
   		for(unsigned int i=1; i<=step_sizes.size();i++){
   			VectorOperations::add_to_std_vector(step_sizes[i-1], direction, temp_pos);
   			positions.push_back(temp_pos);
   		}

    		//Calculate missing L2 values
//		vector<Worker*> workers;
		for(unsigned int i=0; i<n_threads; i++){
			vector<double> new_pos = current_pos;
			VectorOperations::add_to_std_vector(step_size*(i+1), direction, new_pos);
			workers[i]->init(new_pos,seed,i);

   		    workers[i]->start();
   		}
   		for(unsigned int i=0; i<n_threads; i++) workers[i]->join();

   		// Create new seed
   		seed = workers[0]->get_final_configuration();

   		for(unsigned int i=0; i<n_threads; i++){
   			if(i+1<pos_min_displacement) values[i+1] = workers[i]->get_fitness();
   			else values[i+2] = workers[i]->get_fitness();
   			//workers[i]->unlock();
   		}

   		//This should not happen with a distinct minimum
   		if(values[1]>values[0]){
   			cout << "ERROR IN MINIMUM SEARCH: ";
   			cout << "values1-values0=" << values[1]-values[0] <<endl;

   		}

   		//Initialize for next iteration;
   		for(unsigned int i=2; i<values.size(); i++){
   			if(values[i] > values[i-1]){
   				for(unsigned int j=0;j<i-2;j++) before_min_displacement += step_sizes[j];
   				min_displacement = before_min_displacement + step_sizes[i-2];
   				after_min_displacement = min_displacement + step_sizes[i-1];
   				value_at_before_min_displacement = values[i-2];
   				value_at_min_displacement = values[i-1];
   				value_at_after_min_displacement = values[i];
   				ret = orig_start_pos;
   				VectorOperations::add_to_std_vector(min_displacement, direction, ret);
   				break;
   			}
   		}
	}

	for(unsigned int i=0;i<n_threads;i++)
		workers[i]->unlock();

	return ret;
}


void Optimizer::set_gradient_spacing_scaling(vector<double> GRADIENT_SPACING_SCALING) {
	gradient_spacing_scaling = GRADIENT_SPACING_SCALING;
	cout << "Set gradient spacing scaling. First value: " << gradient_spacing_scaling[0] << " --- Last value: " << gradient_spacing_scaling[gradient_spacing_scaling.size()-1] << endl;
}
